from .zmq_utils import init_socket, send_data, recv_data , init_push_socket, init_pull_socket 

class Communicator:
    def __init__(self, is_master, port=5555, handshake_port=6000, result_port=6001):
        self.is_master = is_master
        self.socket = init_socket(bind=is_master, port=port)

        if is_master:
            self.handshake_socket = init_pull_socket(handshake_port)
            self.result_socket = init_pull_socket(result_port)
        else:
            self.handshake_socket = init_push_socket(handshake_port)
            self.result_socket = init_push_socket(result_port)

    def broadcast(self, data):
        send_data(self.socket, data)

    def receive(self):
        return recv_data(self.socket)

    def wait_for_workers(self, num_workers):
        print("[Master] Waiting for workers to connect...")
        ready_workers = 0
        while ready_workers < num_workers:
            msg = self.handshake_socket.recv_string()
            if msg == "READY":
                ready_workers += 1
                print(f"[Master] Received READY from worker {ready_workers}/{num_workers}")
        print("[Master] All workers connected. Proceeding.")

    def notify_master_ready(self):
        self.handshake_socket.send_string("READY")
        print("[Worker] Sent READY to master.")

    def send_result(self, data):
        send_data(self.result_socket, data)

    def receive_result(self):
        return recv_data(self.result_socket)
